#include <stdio.h>
#include <stdlib.h>
#include <time.h>
#include <cuda_runtime.h>

#define BLOCK_SIZE 16

// 全局内存实现的矩阵转置核函数
__global__ void transposeGlobalMemory(float *out, float *in, int n) {
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y * blockDim.y + threadIdx.y;
    
    if (x < n && y < n) {
        out[x * n + y] = in[y * n + x];
    }
}

// 共享内存实现的矩阵转置核函数
__global__ void transposeSharedMemory(float *out, float *in, int n) {
    __shared__ float tile[BLOCK_SIZE][BLOCK_SIZE+1]; // 添加padding避免bank冲突
    
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y * blockDim.y + threadIdx.y;
    
    if (x < n && y < n) {
        tile[threadIdx.y][threadIdx.x] = in[y * n + x];
    }
    
    __syncthreads();
    
    x = blockIdx.y * blockDim.y + threadIdx.x; // 转置后的坐标
    y = blockIdx.x * blockDim.x + threadIdx.y;
    
    if (x < n && y < n) {
        out[y * n + x] = tile[threadIdx.x][threadIdx.y];
    }
}

// 初始化矩阵
void initMatrix(float *mat, int n) {
    for (int i = 0; i < n * n; i++) {
        mat[i] = (float)rand() / RAND_MAX;
    }
}

// 打印矩阵
void printMatrix(float *mat, int n) {
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < n; j++) {
            printf("%.2f ", mat[i * n + j]);
        }
        printf("\n");
    }
}

// 验证转置结果
bool verifyTranspose(float *mat, float *transposed, int n) {
    for (int i = 0; i < n; i++) {
        for (int j = 0; j < n; j++) {
            if (mat[i * n + j] != transposed[j * n + i]) {
                return false;
            }
        }
    }
    return true;
}

int main(int argc, char **argv) {
    if (argc != 2) {
        printf("Usage: %s <matrix_size>\n", argv[0]);
        return 1;
    }
    
    int n = atoi(argv[1]);
    if (n < 512 || n > 2048) {
        printf("Matrix size must be between 512 and 2048\n");
        return 1;
    }
    
    // 分配主机内存
    size_t size = n * n * sizeof(float);
    float *h_A = (float *)malloc(size);
    float *h_AT_global = (float *)malloc(size);
    float *h_AT_shared = (float *)malloc(size);
    
    // 初始化矩阵
    srand(time(NULL));
    initMatrix(h_A, n);
    
    // 分配设备内存
    float *d_A, *d_AT_global, *d_AT_shared;
    cudaMalloc(&d_A, size);
    cudaMalloc(&d_AT_global, size);
    cudaMalloc(&d_AT_shared, size);
    
    // 拷贝数据到设备
    cudaMemcpy(d_A, h_A, size, cudaMemcpyHostToDevice);
    
    // 设置线程块和网格大小
    dim3 block(BLOCK_SIZE, BLOCK_SIZE);
    dim3 grid((n + block.x - 1) / block.x, (n + block.y - 1) / block.y);
    
    // 创建CUDA事件用于计时
    cudaEvent_t start, stop;
    cudaEventCreate(&start);
    cudaEventCreate(&stop);
    float milliseconds = 0;
    
    // 全局内存转置
    cudaEventRecord(start);
    transposeGlobalMemory<<<grid, block>>>(d_AT_global, d_A, n);
    cudaEventRecord(stop);
    cudaEventSynchronize(stop);
    cudaEventElapsedTime(&milliseconds, start, stop);
    
    // 拷贝结果回主机
    cudaMemcpy(h_AT_global, d_AT_global, size, cudaMemcpyDeviceToHost);
    
    printf("Global memory transpose time: %.3f ms\n", milliseconds);
    
    // 共享内存转置
    cudaEventRecord(start);
    transposeSharedMemory<<<grid, block>>>(d_AT_shared, d_A, n);
    cudaEventRecord(stop);
    cudaEventSynchronize(stop);
    cudaEventElapsedTime(&milliseconds, start, stop);
    
    // 拷贝结果回主机
    cudaMemcpy(h_AT_shared, d_AT_shared, size, cudaMemcpyDeviceToHost);
    
    printf("Shared memory transpose time: %.3f ms\n", milliseconds);
    
    // 验证结果
    bool globalCorrect = verifyTranspose(h_A, h_AT_global, n);
    bool sharedCorrect = verifyTranspose(h_A, h_AT_shared, n);
    
    printf("Global memory transpose %s\n", globalCorrect ? "correct" : "incorrect");
    printf("Shared memory transpose %s\n", sharedCorrect ? "correct" : "incorrect");
    
    // 打印部分结果（对于大矩阵，只打印前10x10）
    printf("\nOriginal matrix (first 10x10):\n");
    printMatrix(h_A, 10);
    
    printf("\nGlobal memory transposed (first 10x10):\n");
    printMatrix(h_AT_global, 10);
    
    printf("\nShared memory transposed (first 10x10):\n");
    printMatrix(h_AT_shared, 10);
    
    // 释放资源
    free(h_A);
    free(h_AT_global);
    free(h_AT_shared);
    cudaFree(d_A);
    cudaFree(d_AT_global);
    cudaFree(d_AT_shared);
    cudaEventDestroy(start);
    cudaEventDestroy(stop);
    
    return 0;
}